from langchain_openai import ChatOpenAI
from globle import envs
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import JsonOutputParser,StrOutputParser


'''
需求：利用langchain随机生成3篇知识点长说明，无需排版，段落间换行即可。并保存成txt，字数在1024个token左右

宏观步骤：

1. 先生成3个知识点的名字
2. 根据知识点名分别生成3篇知识点并保存
'''


# 创建模型
model = ChatOpenAI(
    model="deepseek-chat",  # 或具体模型名
    base_url="https://api.deepseek.com",  # DeepSeek API 地址,
)


def get_names(n):
    gen_name_prompt = PromptTemplate.from_template(
        '''随便生成{n}个关于LLM的知识点,返回json。字段
        - name_list 知识点的列表
        '''
    )
    name_chain = gen_name_prompt|model|JsonOutputParser()
    return name_chain.invoke({"n": n})

def gen_essay(name):
    gen_essay_prompt = PromptTemplate.from_template(
        '''根据知识点名字 {name}，生成一篇说明文，无需排版，段落间换行即可。长度在1024个token左右'''
    )

    essay_chain = gen_essay_prompt|model.bind(max_tokens=1024)|StrOutputParser()
    return essay_chain.invoke({"name": name})


def save_text(text, name):
    file_path = f"docs/{name}.txt"
    with open(file_path, "w+", encoding="utf-8") as f:
        f.write(text)

def do():
    names = get_names(3)
    for name in names["name_list"]:
        print(f"开始生成{name}的知识点")
        essay = gen_essay(name)
        save_text(essay, name)
        print(f"生成了{name}的知识点")

if __name__ == "__main__":
    do()